import pickle
from typing import List, Dict

# 利用pyelftools, 从elf的符号表中读出列出全部符号（符号名、符号地址、符号尺寸） 再 查找给定地址在哪个符号中
# pip install pyelftools

from elftools.elf.elffile import ELFFile
from elftools.elf.sections import SymbolTableSection, Section
from elftools.elf.descriptions import  _DESCR_ST_INFO_TYPE
"""
    STT_OBJECT='OBJECT',
    STT_FUNC='FUNC'
"""

class Sym:
    def __init__(self,symbol_name: str, symbol_addr: int, symbol_size: int):
        self.symbol_name: str = symbol_name
        self.symbol_addr: int = symbol_addr
        self.symbol_size: int = symbol_size

    def contain_addr(self,addr:int)->bool:
        contain:bool= addr >= self.symbol_addr and addr <= self.symbol_addr+self.symbol_size
        return contain

    def offsetInFunc(self, addr:int)->int:
        offsetInFunc_: int= addr - self.symbol_addr
        return offsetInFunc_

    def __lt__(self, that):
        if that is None:
            return False

        if type(that) == int:
            return self.symbol_addr<that
        if type(that)==Sym:
            return self.symbol_addr<that.symbol_addr
        return None

    def __gt__(self, that):
        if that is None:
            return False

        if type(that) == int:
            return self.symbol_addr>that
        if type(that)==Sym:
            return self.symbol_addr>that.symbol_addr
        return None
    def __str__(self):
        return f"name:{self.symbol_name},addr:{self.symbol_addr:0x},size:{self.symbol_size}"

def multi_generator(*gens):
    for gen in gens:
        yield from gen


#内存缓存
_mem_cache:Dict[str,List[Sym]]={}
# 从elf的符号表中读出列出全部符号（符号名、符号地址、符号尺寸）
def load_sym_ls__cache(elfFilePath:str)->List[Sym]:
    """
    从elf的符号表中读出列出全部符号（符号名、符号地址、符号尺寸）
    :param elfFilePath: elf文件
    :return: 返回该elf文件的全部符号
    """
# fileName="/pubx/build-llvm15/bin/clang-15"

    #内存缓存, 是为了应对外层的循环调用本方法
    syms_cache= _mem_cache.get(elfFilePath,None)
    if syms_cache is not None:
        return syms_cache

    #{文件缓存开始
    # 排序结果文件缓存, 是为了应对 巨大elf文件 中有大量符号, 一次读取、排序 非常耗时, 因此缓存排序结果到文件以提升效率
    #如果有 排序结果文件缓存, 直接读文件并返回, 以加速
    syms_pickle_fn:str=f"{elfFilePath}.pickle"
    from pathlib import Path
    elf_fp:Path=Path(elfFilePath)
    syms_pickle_fp:Path=Path(syms_pickle_fn)
    if syms_pickle_fp.exists() :
        if syms_pickle_fp.stat().st_ctime > elf_fp.stat().st_ctime:
            with open(syms_pickle_fn,"rb") as pickleRF:
                syms=pickle.load(pickleRF)
                print(f"using cached elf's symbols_pickle_fn:{syms_pickle_fn}")
                #内存缓存, 是为了应对外层的循环调用本方法
                _mem_cache[elfFilePath]=syms
                return syms
        else:
            #缓存结果已经旧了，删掉
            print(f"remove old cached elf's symbols_pickle_fn:{syms_pickle_fn}")
            syms_pickle_fp.unlink()

    syms:List[Sym]=__load_sym_ls(elfFilePath)

    #拿到排序结果后，写一份到文件，下一次不用再排序，以加速
    syms_pickle_fp.unlink(missing_ok=True)
    with open(syms_pickle_fn,"wb") as pickeF:
        pickle.dump(syms,pickeF)
    #文件缓存 结束}



    #内存缓存, 是为了应对外层的循环调用本方法
    _mem_cache[elfFilePath]=syms
    return syms




# 从elf的符号表中读出列出全部符号（符号名、符号地址、符号尺寸）
def __load_sym_ls(elfFilePath:str)->List[Sym]:
    """
    从elf的符号表中读出列出全部符号（符号名、符号地址、符号尺寸）
    :param elfFilePath: elf文件
    :return: 返回该elf文件的全部符号
    """
# fileName="/pubx/build-llvm15/bin/clang-15"
    print(f"load elf's func symbols :{elfFilePath}")
    syms:List[Sym]=[]
    #1. 读取elf文件
    elffile :ELFFile= ELFFile.load_from_path(elfFilePath)

    secLs_dynsym:List[Section]=elffile.iter_sections(type="SHT_SYMTAB")
    secLs_symtab:List[Section]=elffile.iter_sections(type="SHT_DYNSYM")

    sectionK:SymbolTableSection=None

    #2.加载符号表
    for sectionK in multi_generator(secLs_dynsym,secLs_symtab):
        sectionK_num_symbols = sectionK.num_symbols()
        print(f"{sectionK.name}:{sectionK_num_symbols}")
        for symbolJ in sectionK.iter_symbols():
            symbol_name:str=symbolJ.name
            symbol_addr:int=symbolJ.entry["st_value"]
            symbol_size:int=symbolJ.entry["st_size"]
            # print(f'{symbol_name},{symbol_addr:0x},{symbol_size}\n')

            if symbol_size <= 0 or symbolJ.entry['st_info'].type!='STT_FUNC': #STT_OBJECT
                continue
            syms.append( Sym(symbol_name,symbol_addr,symbol_size) )


    #3.对符号按地址排序
    print(f"sort elf's symbols by funcAddr :{elfFilePath}")
    syms=sorted(syms, key=lambda symK:symK.symbol_addr)

    return syms


from bisect import bisect_left,bisect_right

def bi_find_sym(syms:List[Sym], addr:int)->Sym:
    """在已排序的符号表 中 二分查找 给定地址
    :param syms: 符号表
    :param addr: 给定地址
    :return: 在符号表syms中 查找给定地址addr, 返回包含该地址的函数符号, 返回样式为：(symbol_name,symbol_addr,symbol_size)
    """
    # addr:int=0x98ac4a5

    # print(f"addr:{addr:0x}")
    #（二分查找）malloc的调用者地址在syms中哪个函数地址内
    idx=bisect_right(syms, addr, 0, len(syms))
    if idx < 1 or idx >= len(syms):
        return None

    symbolK:Sym=syms[idx-1]
    if symbolK is None:
        return None
    # print(f"symbolK:{symbolK}")
    return symbolK



def findAddr(fileName:str, addr:int):
    # fileName="/pubx/build-llvm15/bin/clang-15"
    # fileName="/lib/x86_64-linux-gnu/libstdc++.so.6"
    # fileName="/tmp/libstdc++.so.6"
    syms:List[Sym]=load_sym_ls__cache(fileName)

    # print("syms loaded")


    # addr: int = 0x98ac4a5
    # addr: int = 0xae9cc
    funcSym:Sym=bi_find_sym(syms, addr)
    print(f"{addr:0x}:{funcSym}")
    # 98ac4a5:('_ZNKSt14default_deleteIKN4llvm8PassInfoEEclEPS2_', 160089202, 59),98ac472

    # addr2: int = 0xa54e641
    # func2:Sym=bi_find_sym(syms, addr2)
    # # a54e641:('_ZN4llvm17deallocate_bufferEPvmm', 173336086, 46),a54e616
    #
    # print(f"{addr2:0x}:{func2}")

    return funcSym

if __name__=="__main__":
    findAddr("/tmp/libstdc++.so.6",0xae9cc)
    #ae9cc:name:_Znwm,addr:ae9b0,size:53
    findAddr("/pubx/build-llvm15/bin/clang-15",0x98ac4a5)
    #98ac4a5:name:_ZNKSt14default_deleteIKN4llvm8PassInfoEEclEPS2_,addr:98ac472,size:59

